# -*- coding:utf-8 -*-
import sqlite3
from prompt_toolkit import PromptSession
from prompt_toolkit.history import FileHistory
from colorama import Fore, Style
from completer import SQLCompleter
from formatter import handle_pagination
from tabulate import tabulate  # 添加 tabulate 导入
import os

class SQLiteCLI:
    def __init__(self, db_name=":memory:"):
        self.conn = sqlite3.connect(db_name)
        self.cursor = self.conn.cursor()
        self.sql_completer = SQLCompleter(self)
        self.prompt_message = "sqlite> "
        
        # 创建历史记录目录
        history_dir = os.path.expanduser('~/.sqlite_cli')
        os.makedirs(history_dir, exist_ok=True)
        history_file = os.path.join(history_dir, 'history.txt')
        
        # 设置带历史记录的会话
        self.session = PromptSession(
            message=self.prompt_message,
            completer=self.sql_completer,
            history=FileHistory(history_file)
        )

    def get_table_names(self):
        try:
            self.cursor.execute("SELECT name FROM sqlite_master WHERE type='table'")
            return [table[0] for table in self.cursor.fetchall()]
        except sqlite3.Error:
            return []

    def show_tables(self):
        try:
            self.cursor.execute("SELECT name FROM sqlite_master WHERE type='table'")
            tables = self.cursor.fetchall()
            if tables:
                print(Fore.CYAN + "\nAvailable tables:" + Style.RESET_ALL)
                for table in tables:
                    print(Fore.GREEN + f"- {table[0]}" + Style.RESET_ALL)
            else:
                print(Fore.YELLOW + "\nNo tables found in database." + Style.RESET_ALL)
        except sqlite3.Error as e:
            print(Fore.RED + f"Error: {e}" + Style.RESET_ALL)

    def show_table_structure(self, table_name):
        try:
            self.cursor.execute(f"PRAGMA table_info({table_name})")
            columns = self.cursor.fetchall()
            if columns:
                # 构建表结构信息
                headers = ['Column', 'Type', 'Nullable', 'Default', 'Primary Key']
                rows = [[col[1], col[2], 'NOT NULL' if not col[3] else 'NULL',
                        col[4] if col[4] is not None else 'NULL',
                        'YES' if col[5] else 'NO'] for col in columns]
                print(Fore.CYAN + f"\nTable structure for '{table_name}':" + Style.RESET_ALL)
                print(Fore.CYAN + tabulate(rows, headers=headers, tablefmt='grid') + Style.RESET_ALL)
            else:
                print(Fore.YELLOW + f"\nTable '{table_name}' not found." + Style.RESET_ALL)
        except sqlite3.Error as e:
            print(Fore.RED + f"Error: {e}" + Style.RESET_ALL)

    def execute_query(self, query):
        try:
            query = query.strip().rstrip(';')
            
            # 处理 DESC 命令
            if query.upper().startswith('DESC '):
                table_name = query[5:].strip()
                self.show_table_structure(table_name)
                return
            
            if query.upper() == 'SHOW TABLES':
                self.show_tables()
                return

            self.cursor.execute(query)
            if query.strip().upper().startswith(('SELECT', 'PRAGMA', 'SHOW')):
                rows = self.cursor.fetchall()
                if rows:
                    headers = [description[0] for description in self.cursor.description]
                    handle_pagination(rows, headers)
                else:
                    print(Fore.YELLOW + "Query returned no results." + Style.RESET_ALL)
            else:
                self.conn.commit()
                print(Fore.GREEN + "Query executed successfully." + Style.RESET_ALL)
        except sqlite3.Error as e:
            print(Fore.RED + f"Error: {e}" + Style.RESET_ALL)

    def run(self):
        print(Fore.CYAN + "Welcome to SQLite CLI! Type 'exit' or 'quit' to exit." + Style.RESET_ALL)
        print(Fore.CYAN + "Type 'show tables' to list all tables." + Style.RESET_ALL)
        while True:
            try:
                query = self.session.prompt()
                if query.lower() in ('exit', 'quit'):
                    break
                if query.strip():
                    self.execute_query(query)
            except KeyboardInterrupt:
                continue
            except EOFError:
                break
        self.conn.close()
        print(Fore.CYAN + "\nGoodbye!" + Style.RESET_ALL)